from aiohttp import web
import asyncio
import time
import multiprocessing
import argparse
from datetime import datetime

connected_clients = set()
message_counter = 0
message_history = []

# 创建一个队列用于进程间通信
message_queue = multiprocessing.Queue()

async def websocket_handler(request):
    global message_counter
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    connected_clients.add(ws)

    # 发送历史消息给新连接的用户
    for msg in message_history:
        await ws.send_str(msg)

    try:
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                if msg.data == 'ping':
                    await ws.send_str('pong')
                else:
                    message_counter += 1
                    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    formatted_msg = f"{message_counter}: {timestamp}: {msg.data}"
                    message_history.append(formatted_msg)
                    # 广播消息给所有连接的客户端，包括发送者自己
                    for client in connected_clients:
                        await client.send_str(formatted_msg)
            elif msg.type == web.WSMsgType.ERROR:
                print('ws connection closed with exception %s' % ws.exception())
    finally:
        connected_clients.remove(ws)
        await ws.close()

async def index(request):
    # 获取命令行参数中的主机和端口
    host = request.app['host']
    port = request.app['port']

    content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>WebSocket Server</title>
        <style>
            table {{
                width: 100%;
                border-collapse: collapse;
                table-layout: fixed; /* 固定表格布局 */
            }}
            table, th, td {{
                border: 1px solid black;
            }}
            th, td {{
                padding: 10px;
            }}
            th {{
                position: sticky;
                top: 0;
                background-color: #f1f1f1;
            }}
            th:first-child, td:first-child {{
                width: 50px; /* 序号列的宽度 */
                text-align: center; /* 序号居中显示 */
            }}
            th:nth-child(2), td:nth-child(2) {{
                width: 150px; /* 时间列的宽度 */
                text-align: center; /* 时间居中显示 */
                white-space: nowrap; /* 防止时间换行 */
            }}
            td:last-child {{
                width: auto; /* 文字列的宽度 */
                text-align: left; /* 消息左对齐显示 */
                white-space: nowrap; /* 防止文字换行 */
                overflow: hidden; /* 隐藏溢出的文字 */
                text-overflow: ellipsis; /* 显示省略号 */
            }}
        </style>
    </head>
    <body>
        <h2>WebSocket Server</h2>
        <table id="messages">
            <thead>
                <tr>
                    <th>序号</th>
                    <th>时间</th>
                    <th>消息</th>
                </tr>
            </thead>
            <tbody id="messages-body">
            </tbody>
        </table>
        <input type="text" id="message_input" />
        <button onclick="sendMessage()">Send</button>
        <script>
            var ws = new WebSocket("ws://{host}:{port}/ws");
            var pingInterval;
            var reconnectAttempts = 0;
            var maxReconnectAttempts = 5;

            function connectWebSocket() {{
                ws = new WebSocket("ws://{host}:{port}/ws");
                ws.onopen = function() {{
                    console.log("WebSocket connection established.");
                    reconnectAttempts = 0;
                    pingInterval = setInterval(function() {{
                        ws.send("ping");
                    }}, 5000); // 每5秒发送一次心跳
                }};

                ws.onmessage = function(event) {{
                    if (event.data === 'pong') {{
                        return; // 忽略心跳响应
                    }}
                    var messageElement = document.createElement("tr");
                    var parts = event.data.split(': ');
                    var index = parts[0];
                    var timestamp = parts[1];
                    var message = parts[2];
                    messageElement.innerHTML = `<td>${{index}}</td><td>${{timestamp}}</td><td>${{message}}</td>`;
                    document.getElementById("messages-body").prepend(messageElement);
                }};

                ws.onclose = function(event) {{
                    console.log("WebSocket connection closed. Reconnecting...");
                    clearInterval(pingInterval);
                    if (reconnectAttempts < maxReconnectAttempts) {{
                        reconnectAttempts++;
                        setTimeout(connectWebSocket, 2000); // 2秒后尝试重新连接
                    }} else {{
                        console.log("Max reconnect attempts reached. Please reload the page.");
                    }}
                }};
            }}

            connectWebSocket();

            function sendMessage() {{
                var message = document.getElementById("message_input").value;
                ws.send(message);
                document.getElementById("message_input").value = '';
            }}
        </script>
    </body>
    </html>
    """
    return web.Response(text=content, content_type='text/html')

async def start_server(host, port):
    app = web.Application()
    app['host'] = host
    app['port'] = port
    app.router.add_get('/', index)
    app.router.add_get('/ws', websocket_handler)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, host, port)
    await site.start()
    print(f"Server started at http://{host}:{port}")

    # 启动消息处理任务
    asyncio.create_task(process_messages())

async def process_messages():
    global message_counter
    while True:
        if not message_queue.empty():
            message = message_queue.get()
            message_counter += 1
            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            formatted_msg = f"{message_counter}: {timestamp}: {message}"
            message_history.append(formatted_msg)
            # 广播消息给所有连接的客户端
            for client in connected_clients:
                await client.send_str(formatted_msg)
        await asyncio.sleep(1)  # 每秒检查一次队列

def message_producer(queue):
    # 模拟从其他地方产生消息
    import random
    import time
    while True:
        message = f"External message {random.randint(1, 100)}"
        queue.put(message)
        time.sleep(5)  # 每5秒产生一条消息

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='WebSocket Server with Command Line Port')
    parser.add_argument('--host', type=str, default='127.0.0.1', help='Host to run the server on')
    parser.add_argument('--port', type=int, default=9000, help='Port number to run the server on')
    args = parser.parse_args()

    loop = asyncio.get_event_loop()

    # 启动消息生产者进程
    producer_process = multiprocessing.Process(target=message_producer, args=(message_queue,))
    producer_process.start()

    loop.run_until_complete(start_server(args.host, args.port))
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass
    finally:
        producer_process.terminate()
        loop.close()
